import pandas as pd
import os
'''
# 定义标签映射函数
def get_label(filename):
    if filename.startswith('E3TF普通'):
        return 'E3TF_0'  # E3TF普通 对应 E3TF_0
    elif filename.startswith('E3TF'):
        return 'E3TF_1'  # ESTF 对应 E3TF_1
    elif filename.startswith('P8普通'):
        return 'P8_0'  # P8普通 对应 P8_0
    elif filename.startswith('P8椭圆'):
        return 'P8_1'  # P8椭圆 对应 P8_1
    elif filename.startswith('P8跑道'):
        return 'P8_2'  # P8跑道 对应 P8_2
    elif filename.startswith('R135普通'):
        return 'R135_0'  # R135普通 对应 R135_0
    elif filename.startswith('R135'):
        return 'R135_1'  # R135 对应 R135_1
    else:
        return None  # 无法匹配的文件名
'''


def get_label(filename):
    if filename.startswith('E3TF普通'):
        return 'PT'  # E3TF普通 对应 PT
    elif filename.startswith('E3TF'):
        return 'E3TF_1'  # E3TF 对应 E3TF_1
    elif filename.startswith('P8普通'):
        return 'PT'  # P8普通 对应 PT
    elif filename.startswith('P8椭圆'):
        return 'P8_1'  # P8椭圆 对应 P8_1
    elif filename.startswith('P8跑道'):
        return 'P8_2'  # P8跑道 对应 P8_2
    elif filename.startswith('R135普通'):
        return 'PT'  # R135普通 对应 PT
    elif filename.startswith('R135'):
        return 'R135_1'  # R135 对应 R135_1
    else:
        return None  # 无法匹配的文件名
# 定义目录
input_directory = r'D:\pythonProject\模式识别1115\测试集_备份'  # 替换为您的目录路径

# 遍历目录下的所有 CSV 文件
for filename in os.listdir(input_directory):
    if filename.endswith('.csv'):
        # 读取 CSV 文件
        file_path = os.path.join(input_directory, filename)
        df = pd.read_csv(file_path)

        # 获取对应标签
        label = get_label(filename)
        if label is not None:
            # 添加标签列
            df['Label'] = label

            # 保存回 CSV 文件
            df.to_csv(file_path, index=False)
            print(f'Updated {filename} with label {label}.')
        else:
            print(f'No match for {filename}.')